#!/usr/bin/perl -w
use strict;
use DBI;
use DBD::Pg;
use Getopt::Long;
use Digest::MD4;

# constants
use constant RANK => 0;
use constant ID => 1;
use constant PARENT => 2;
use constant DIVISION => 3;

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#	usage

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
my $usage = <<'USAGE';

USAGE:

    panda_load_taxa.pl --nodes_file <FILE> --names_file <FILE> --gi_file <FILE> --dbname <FILE>

    Mandatory arguments:

         --nodes_file FILE   nodes.dmp from NCBI containing nodes from the taxid tree
         --names_file FILE   names.dmp from NCBI containing names for each taxid
         --gi_file    FILE   gi_taxid_prot.dmp from NCBI for gi => taxid lookups
         --dbname     FILE   the postgreSQL database


    Optional arguments:

        [--user STRING]      the postgreSQL user name
        [--host STRING]      the postgreSQL database server
        [--password STRING]  the postgreSQL database password
        [--verbose]

        [--help]


RESULTS:

    This script loads the taxonomy tree from NCBI (in a flattened form) into postgreSQL.

USAGE



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#	Get options

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


# options variables
my $user			= 'postgres';
my $password		= '';
my $host			= '';

my $help			= '';
my $dbname			= '';
my $verbose;


my $nodes_file ;
my $names_file;
my $gi_file   ;


GetOptions ('nodes_file=s'	=> \$nodes_file,
			'names_file=s'	=> \$names_file,
			'gi_file=s'		=> \$gi_file,
			'help'			=> \$help,
			'verbose'		=> \$verbose,
			'user=s'		=> \$user,
			'password'		=> \$password,
			'host=s'		=> \$host,
			'dbname=s'		=> \$dbname,
			);



die $usage if ($help);
die "\nERROR:\n\tMissing Arguments:\n" .$usage
				unless ($nodes_file  &&
						$names_file &&
						$gi_file    &&
						$dbname);
				
				
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#	Initialize

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
# open files
open NODES, $nodes_file  or die "$nodes_file  not found";
open NAMES, $names_file or die "$names_file not found";
open GIS, $gi_file    or die "$gi_file    not found";

# connect to database
my $dbihost = ($host) ? ";host=$host" : "";
my $dbh = DBI->connect( "dbi:Pg:dbname=$dbname$dbihost",
                        $user,
                        $password,
						{
                          RaiseError => 0,
                          AutoCommit => 1,
						  PrintError => 0
                        }
                      )
		or die "Database connection to $dbname failed:\n$DBI::errstr\n";

my $md4 = new Digest::MD4;



eval
{
	$dbh->do("drop table newtaxid_from_taxname");
};
eval
{
	$dbh->do("drop table newtaxa");
};
eval
{
	$dbh->do("drop table newtaxid_from_gi");
};
$dbh->{AutoCommit} = 0;
$dbh->{PrintError} = 1;


#****************************************************************************************

# get list of desired taxonomical ranks

#****************************************************************************************
my $CNT_RANKS = 0;
my %ranks;
{
	while (<DATA>)
	{
		chomp;
		$ranks{$_} = ++$CNT_RANKS;
	}
}




#****************************************************************************************

#	Read in all taxonomical names

#****************************************************************************************
print STDERR "Entering taxonomical names to NCBI taxid lookup data\n" if ($verbose);
my %names;
	
	eval
	{
	

		# prepare SQL
		$dbh->do("Create table newtaxid_from_taxname( tax_name_hash OID[4], tax_id  OID)");
		my $sth = $dbh->prepare("COPY newtaxid_from_taxname ( tax_name_hash, tax_id) FROM STDIN");
		$sth->execute();
	
		print STDERR "\t[" if ($verbose);
		my $cnt_names = 0;
		while (<NAMES>)
		{
			# put name into panda
			next unless /^(\d+)\t\|\t([^\t]+)\t\|/;
			if (++$cnt_names % 4000 == 0)
			{
				if ($verbose)
				{
					print STDERR "\n\t" if (($cnt_names / 4000) % 50 == 0);
					print STDERR '.';
				}
			}
			my @digest = unpack("L*", $md4->hash($2));
			my $line = '{'.join(',', @digest)."}\t$1\n";
			$dbh->func($line, 'putline');
			
			# save name for lookup tree later
			#1	|	root	|		|	scientific name	|
			next unless /^(\d+)\t\|\t([^\t]+)\t\|[^\|]+\|\tscientific name/;
			$names{$1} = $2;
		}
		$dbh->func("\\.\n", 'putline');
		$dbh->func('endcopy');
		print STDERR "]\n"  if ($verbose);
		print STDERR "\t$cnt_names\tName / taxid pairs entered into postgreSQL\n"  if ($verbose);
	};
	if( $@ )
	{
		warn "Database error: $DBI::errstr\n$@\n";
		$dbh->rollback(); #just die if rollback is failing
		exit;
	}

#****************************************************************************************

#	Read in gis

#****************************************************************************************
print STDERR "Entering NCBI gi to NCBI taxid lookup data\n" if ($verbose);
	
	eval
	{
	
		# prepare SQL
		$dbh->do("Create table newtaxid_from_gi( gi INTEGER, tax_id  OID)");
		my $sth = $dbh->prepare("COPY newtaxid_from_gi FROM STDIN");
		$sth->execute();
		
		print STDERR "\t[" if ($verbose);
		my $cnt_gis = 0;
		while (<GIS>)
		{
			# put name into panda
			next unless /^(\d+\t\d+\n)/;
			if (++$cnt_gis % 30000 == 0)
			{
				if ($verbose)
				{
					print STDERR "\n\t" if (($cnt_gis / 30000) % 50 == 0);
					print STDERR '.';
				}
			}
			$dbh->func($1, 'putline');
		}
		
		$dbh->func("\\.\n", 'putline');
		$dbh->func('endcopy');
		print STDERR "]\n"  if ($verbose);
		print STDERR "\t$cnt_gis\tGI / taxid pairs entered into postgreSQL\n" if ($verbose);
	};
	if( $@ )
	{
		warn "Database error: $DBI::errstr\n$@\n";
		$dbh->rollback(); #just die if rollback is failing
		exit;
	}








#****************************************************************************************

#	Read in all tax tree in child->Parent structure

#****************************************************************************************
# read in tax ids and parents
print STDERR "Get list of TAXIDs \n"  if ($verbose);
	my %nodes;
	while (<NODES>)
	{
		next unless /	(\d+)\t\|				# id
						\t(\d+)\t\|				# parent id
						\t([^\t]+)\t\|			# rank e.g. species or family
						[^\|]+\|				# ignore embl code
						\t(\d+)\t				# division code
						/x;
		my $rank =  (exists $ranks{$3}) ? $ranks{$3} : 0;
		$nodes{$1} = [$rank, $1, $2, $4];
	}
	print STDERR "Finishing getting @{[scalar keys %nodes]} list of TAXIDs \n" if ($verbose);







#****************************************************************************************

# 	Join up nodes to make tree

#****************************************************************************************

foreach my $id (keys %nodes)
{
	my $node = $nodes{$id};
	die "parent id ($node->[PARENT]) for taxid $id not found!"
									unless exists $nodes{$node->[PARENT]};
	$node->[PARENT] = $nodes{$node->[PARENT]};
}


#****************************************************************************************

#	Walk up each node to top and output flattened tree

#****************************************************************************************
my $cnt_subspecies_nodes = 0;
eval
{

# prepare SQL
$dbh->do("Create table newtaxa( tax_id OID,".
					  "tax_name TEXT,".
					  "division INTEGER,".
					  "species OID,".
					  "genus OID,".
					  "family OID,".
					  "tax_order OID,".
					  "class OID,".
					  "phylum OID,".
					  "kingdom OID,".
					  "superkingdom OID)");
my $sql = "INSERT INTO newTAXA VALUES ( " . ("?,") x ($CNT_RANKS+2) . "? )";
my $sth = $dbh->prepare( $sql );

	print STDERR "\t["  if ($verbose);
	foreach my $id (keys %nodes)
	{
		my $curr_node = $nodes{$id};
		my @tree = ($id, $names{$id}, $curr_node->[DIVISION], (0) x $CNT_RANKS);
		while ($curr_node->[PARENT] != $curr_node)
		{
			# have the entry for the appropriate rank point to myself
			$tree[$curr_node->[RANK]+2] = $curr_node->[ID]
									if ($curr_node->[RANK]);
			
			# move up to parent node
			$curr_node = $curr_node->[PARENT];
		}

	#	print (join "\t", @tree);
	#	print "\n";

		$sth->execute(@tree);

		if (++$cnt_subspecies_nodes % 2000 == 0)
		{
			if ($verbose)
			{
				print STDERR "\n\t" if (($cnt_subspecies_nodes / 2000) % 50 == 0);
				print STDERR '.';
			}
		}
		

	}
	print STDERR "]\n"  if ($verbose);
	$sth->finish();
	print STDERR "Nodes = $cnt_subspecies_nodes\n"  if ($verbose);
};
if( $@ )
{
	warn "Database error: $DBI::errstr\n$@\n";
	$dbh->rollback(); #just die if rollback is failing
	exit;
}
#****************************************************************************************

#	Cleanup

#****************************************************************************************
	$dbh->commit();
	$dbh->{AutoCommit} = 1;
	$dbh->{RaiseError} = 0;
	$dbh->{PrintError} = 0;
	$dbh->do("DROP TABLE taxa");
	$dbh->do("DROP INDEX i_taxa__tax_id");
	$dbh->do("DROP INDEX i_taxa__name");
	$dbh->do("DROP INDEX i_taxa__division");
	$dbh->do("DROP INDEX i_taxa__species");
	$dbh->{PrintError} = 1;
	$dbh->do("ALTER TABLE newtaxa RENAME TO taxa");
	$dbh->do("CREATE UNIQUE INDEX i_taxa__tax_id ON taxa (tax_id)");
	$dbh->do("CREATE INDEX i_taxa__name ON taxa (tax_name)");
	$dbh->do("CREATE INDEX i_taxa__division ON taxa (division)");
	$dbh->do("CREATE INDEX i_taxa__species ON taxa (species)");
	$dbh->do("DROP TABLE taxid_from_taxname");
	$dbh->do("DROP TABLE taxid_from_gi");
	$dbh->do("ALTER TABLE newtaxid_from_taxname RENAME TO taxid_from_taxname");
	$dbh->do("ALTER TABLE newtaxid_from_gi RENAME TO taxid_from_gi");
$dbh->disconnect();

print STDERR "Success!\n"  if ($verbose);
__DATA__
species
genus
family
order
class
phylum
kingdom
superkingdom
